# -*- coding: utf-8 -*-
#!/usr/bin python

# soft-NMS meachism

def soft_nms(boxes, sigma, Nt, thresh, method):
    """
    boxes: np.array([[x1, y1, x2, y2, socre], [....], [....]])
    sigma: soft_nma parameter
    Nt: iou threshold
    thresh: ?
    method: soft-NMS method
    """
    N = boxes.shape[0]
    for i in range(N):
        max_score = boxes[i, 4]
        max_pos = i

        tx1 = boxes[i, 0]
        ty1 = boxes[i, 1]
        tx2 = boxes[i, 2]
        ty2 = boxes[i, 3]
        tscore = boxes[i, 4]

        pos  = i + 1
        # get max_score box
        while pos < N:
            if max_score < boxes[pos, 4]:
                max_score = boxes[pos, 4]
                max_pos = pos
            pos = pos + 1

        # add max_score box as a detection
        boxes[i, 0] = boxes[max_pos, 0]
        boxes[i, 1] = boxes[max_pos, 1]
        boxes[i, 2] = boxes[max_pos, 2]
        boxes[i, 3] = boxes[max_pos, 3]
        boxes[i, 4] = boxes[max_pos, 4]

        # swap i-th box with position of max_score box
        boxes[max_pos, 0] = tx1
        boxes[max_pos, 1] = ty1
        boxes[max_pos, 2] = tx2
        boxes[max_pos, 3] = ty2
        boxes[max_pos, 4] = tscore

        tx1 = bboxes[i, 0]
        ty1 = bboxes[i, 1]
        tx2 = bboxes[i, 2]
        ty2 = bboxes[i, 3]
        tscore = bboxes[i, 4]

        pos = i + 1
        # NMS iterations, note that N changes if detection boxes fall below
        # threshold
        while pos < N:
            x1 = bboxes[pos, 0]
            y1 = bboxes[pos, 1]
            x2 = bboxes[pos, 2]
            y2 = bboxes[pos, 3]
            score = boxes[pos, 4]

            area = (x2 - x1 + 1) * ( y2 - y1 + 1)
            iw = (min(tx2, x2) - max(tx1, x1) + 1)
            if iw > 0:
                ih = (min(ty2, y2) - max(ty1, y1) + 1)
                if ih > 0:
                    ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw *ih)
                    ov = iw*ih / ua # iou between max_score box and detection box
                    
                    if method == 1:
                        if ov > Nt:
                            weight = 1 - ov
                        else:
                            weight = 1
                    elif method == 2:
                        weight = np.exp(-(ov * ov) / sigma)
                    else:
                        if ov > Nt:
                            weigth = 0
                        else:
                            weight = 1

                    boxes[pos, 4] = weigth * boxes[pos, 4]

                    # if box score falls below threshold, discard the box by
                    # swapping with last box update N
                    if boxes[pos, 4] < threshold:
                        boxes[pos, 0] = boxes[N-1, 0]
                        boxes[pos, 1] = boxes[N-1, 1]
                        boxes[pos, 2] = boxes[N-1, 2]
                        boxes[pos, 3] = boxes[N-1, 3]
                        boxes[pos, 4] = boxes[N-1, 4]
                        N = N - 1
                        pos = pos -1
            pos = pos + 1

        keep = [i for i in range(N)]
        return keep
